package com.example.mplayer.player

import androidx.media3.datasource.DataSource
import androidx.media3.datasource.DataSpec
import androidx.media3.datasource.TransferListener
import android.net.Uri
import androidx.media3.common.util.UnstableApi
import java.security.PrivateKey
import javax.crypto.Cipher

@UnstableApi
class EncryptionKeyDataSource(
    private val upstream: DataSource,
    private val privateKey: PrivateKey
) : DataSource {

    override fun addTransferListener(transferListener: TransferListener) {
        upstream.addTransferListener(transferListener)
    }

    override fun open(dataSpec: DataSpec): Long {
        return upstream.open(dataSpec)
    }

    override fun read(buffer: ByteArray, offset: Int, length: Int): Int {
        val bytesRead = upstream.read(buffer, offset, length)
        if (bytesRead > 0 && uri?.path?.endsWith("/aes.key") == true) {
            // Decrypt the AES key using RSA private key
            val cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding")
            cipher.init(Cipher.DECRYPT_MODE, privateKey)
            val decryptedKey = cipher.doFinal(buffer.copyOfRange(offset, offset + bytesRead))
            System.arraycopy(decryptedKey, 0, buffer, offset, decryptedKey.size)
            return decryptedKey.size
        }
        return bytesRead
    }

    override fun getUri(): Uri? = upstream.uri

    override fun close() {
        upstream.close()
    }
}